from ultralytics import YOLO

model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)

def main():
    results = model.train(
        data='./CreateDataset/dataset1/data.yaml',  # 数据集配置文件路径
        epochs=60,                 # 训练总轮次
        patience=10,                # 早停等待轮次
        batch=16,                   # 批次大小（根据GPU显存调整）
        imgsz=640,                  # 输入图像尺寸
        save=True,                  # 保存训练结果
        save_period=10,             # 每10轮保存一次检查点
        cache=False,               # 禁用数据缓存（显存小时建议False）
        device=0,                # 使用GPU设备（'0'或'cpu'）
        workers=8,                  # 数据加载线程数
        project='runs/train',       # 结果保存目录
        name='exp',                 # 实验名称
        optimizer='auto',          # 自动选择优化器
        lr0=0.01,                  # 初始学习率
        lrf=0.01,                  # 最终学习率 = lr0 * lrf
        momentum=0.937,             # SGD动量
        weight_decay=0.0005,        # 权重衰减系数
        warmup_epochs=3.0,          # 热身轮次
        box=7.5,                    # box损失权重
        cls=0.5,                    # 分类损失权重
        dfl=1.5,                    # dfl损失权重
        verbose=True,               # 打印详细输出
    )

if __name__ == '__main__':
    main()